# coding:utf8
from config import opt
import os
import torch as t
import models
from data.dataset import AVD
from torch.tensor import Tensor
from torch.nn import functional
from torch.utils.data import DataLoader
from torch.autograd import Variable
from utils.visualize import Visualizer
from tqdm import tqdm


def test(**kwargs):
    opt.parse(kwargs)
    # configure model
    model = getattr(models, opt.model)().eval()
    if opt.load_model_path:
        model.load(opt.load_model_path)
    if opt.use_gpu: model.cuda()

    # data
    train_data = AVD(opt.test_data_root, test=True)
    test_dataloader = DataLoader(train_data, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers)
    results = []
    for ii, (data, path) in enumerate(test_dataloader):
        input = Tensor(data, volatile=True)
        if opt.use_gpu: input = input.cuda()
        score = model(input)
        probability = functional.softmax(score)[:, 0].data.tolist()
        # label = score.max(dim = 1)[1].data.tolist()

        batch_results = [(path_, probability_) for path_, probability_ in zip(path, probability)]

        results += batch_results
    write_csv(results, opt.result_file)

    return results


def write_csv(results, file_name):
    import csv
    with open(file_name, 'w') as f:
        writer = csv.writer(f)
        writer.writerow(['id', 'label'])
        writer.writerows(results)


def train(**kwargs):
    opt.parse(kwargs)
    vis = Visualizer(opt.env)

    # step1: configure model
    model = getattr(models, opt.model)()
    if opt.load_model_path:
        model.load(opt.load_model_path)
    if opt.use_gpu: model.cuda()

    # step2: data
    train_data = AVD(opt.train_data_root, train=True)
    val_data = AVD(opt.train_data_root, train=False)
    train_dataloader = DataLoader(train_data, opt.batch_size,
                                  shuffle=True, num_workers=opt.num_workers)
    val_dataloader = DataLoader(val_data, opt.batch_size,
                                shuffle=False, num_workers=opt.num_workers)

    # step3: criterion and optimizer
    criterion = t.nn.CrossEntropyLoss()
    lr = opt.lr
    optimizer = t.optim.Adam(model.parameters(), lr=lr, weight_decay=opt.weight_decay)

    # step4: meters
    # loss_meter = meter.AverageValueMeter()
    # confusion_matrix = meter.ConfusionMeter(2)
    previous_loss = 1e100

    # train
    for epoch in range(opt.max_epoch):

        # loss_meter.reset()
        # confusion_matrix.reset()

        for ii, (data, label) in tqdm(enumerate(train_dataloader), total=len(train_data)):

            # train model 
            input = Tensor(data)
            target = Tensor(label)
            if opt.use_gpu:
                input = input.cuda()
                target = target.cuda()

            optimizer.zero_grad()
            score = model(input)
            loss = criterion(score, target)
            loss.backward()
            optimizer.step()
            if ii % opt.print_freq == opt.print_freq - 1:
                if os.path.exists(opt.debug_file):
                    pass

        model.save()



def val(model, dataloader):

    model.eval()
    # confusion_matrix = meter.ConfusionMeter(2)
    for ii, data in enumerate(dataloader):
        input, label = data
        val_input = Tensor(input, volatile=True)
        val_label = Tensor(label.type(t.LongTensor), volatile=True)
        if opt.use_gpu:
            val_input = val_input.cuda()
            val_label = val_label.cuda()
        score = model(val_input)
        # confusion_matrix.add(score.data.squeeze(), label.type(t.LongTensor))

    model.train()
    # cm_value = confusion_matrix.value()
    # accuracy = 100. * (cm_value[0][0] + cm_value[1][1]) / (cm_value.sum())
    # return confusion_matrix, accuracy


def help():
    """
        Print help information: python file.py help
    """

    print('''
    usage : python file.py <function> [--args=value]
    <function> := train | test | help
    example: 
            python {0} train --env='env0701' --lr=0.01
            python {0} test --dataset='path/to/dataset/root/'
            python {0} help
    avaiable args:'''.format(__file__))

    from inspect import getsource
    source = (getsource(opt.__class__))
    print(source)


if __name__ == '__main__':
    import fire

    fire.Fire()
